from word_parser.word_parser import WordParser
from word_parser.word_parser_two import WordParserTwo
import argparse
import os, sys
from time import time

if __name__ == '__main__':
    argparser = argparse.ArgumentParser()
    argparser.add_argument('input', action='store', type=str, help='Input file path')
    argparser.add_argument('output', action='store', nargs='?', type=str, default='output.txt', help='Output file path')
    argparser.add_argument('-m','--model', action='store', type=int, choices = [2,3], default=3, help='Use Bi-gram or Tri-gram')
    argparser.add_argument('-t', '--train', action='store', help='Train the model or not')
    argparser.add_argument('--gamma',action='store',type=float)
    argparser.add_argument('--theta',action='store',type=float)
    args = argparser.parse_args()

    current_word_dir = os.path.abspath('.')
    if args.model == 3:
        word_parser = WordParser()
        save_load_prefix = 'train_three'
    else:
        word_parser = WordParserTwo()
        save_load_prefix = 'train'

    if args.train:
        word_parser.init_freq(os.path.join(current_word_dir,'resource', 'hanzi_table.txt'))
        resource_path = args.train
        # 训练模型
        for file_path in os.listdir(resource_path):
            word_parser.statitical_frequency(os.path.join(resource_path, file_path))
            print('Train {} success!'.format(file_path))
        print('Start to save...\r', end='')
        word_parser.save('{}_data.cache'.format(save_load_prefix))
        print('Save success!')
    else:
        # 读入保存好的模型
        print('Loading cache...\r', end='')
        word_parser.load('{}_data.cache'.format(save_load_prefix))
        print('Load cache success!')

    word_parser.summary()

    handile_line = 0
    ans = []
    with open(args.input, 'r', encoding='utf-8') as f:
       line = f.readline().strip('\n').lower()
       while line:
            handile_line += 1
            start = time()
            res = word_parser.find_match(line)
            end = time()
            ans.append(res)
            print('处理第{}条拼音\t 耗时:{:.2f} s\r'.format(handile_line, end-start),end='')
            line = f.readline().strip('\n').lower()
    with open(args.output,'w',encoding='utf-8') as f:
        for res in ans:
            f.write(res+'\n')
    